from qiskit import QuantumCircuit
from typing import List, Optional, Union


class GHZCircuit:
    """
    A class for creating Greenberger-Horne-Zeilinger (GHZ) state circuits.
    
    The GHZ state is a maximally entangled quantum state of the form:
    |GHZ⟩ = (|00...0⟩ + |11...1⟩)/√2
    
    This class provides methods to create and customize GHZ state circuits
    for various quantum computing applications and benchmarks.
    """
    
    def __init__(self, num_qubits: int, qubit_indices: Optional[List[int]] = None):
        """
        Initialize the GHZ circuit generator.
        
        Args:
            num_qubits: Total number of qubits in the circuit
            qubit_indices: Specific qubits to use for the GHZ state. If None, all qubits are used.
        """
        self.num_qubits = num_qubits
        
        # If qubit_indices is not provided, use all qubits
        if qubit_indices is None:
            self.qubit_indices = list(range(num_qubits))
        else:
            self.qubit_indices = qubit_indices
            
        self.circuit = None
    
    def generate_circuit(self, measure: bool = True, qed: bool = False, pilocy: str = None) -> QuantumCircuit:
        """
        Generate a quantum circuit that prepares a GHZ state.
        
        Args:
            measure: Whether to include measurement operations in the circuit
            qed: Whether to use qed techique, default is False
            pilocy: The pilocy of the qed mapping, default is None

        Returns:
            A quantum circuit that prepares a GHZ state
        """
        # Create a new quantum circuit
        qc = QuantumCircuit(self.num_qubits, len(self.qubit_indices) if measure else 0)
        
        # Apply Hadamard to the first qubit
        qc.h(self.qubit_indices[0])
        
        # Apply CNOT gates to entangle all qubits
        for i in range(len(self.qubit_indices) - 1):
            control = self.qubit_indices[i]
            target = self.qubit_indices[i + 1]
            qc.cx(control, target)

            if qed:
                if pilocy == "01&10":
                    qc.x(target)


        # Add measurements if requested
        if measure:
            for i, qubit_idx in enumerate(self.qubit_indices):
                qc.measure(qubit_idx, i)
        
        self.circuit = qc
        return qc
    
    def generate_circuit_with_phase(self, phase: float, measure: bool = True) -> QuantumCircuit:
        """
        Generate a quantum circuit that prepares a GHZ state with a relative phase.
        
        The state will be of the form: |GHZ⟩ = (|00...0⟩ + e^(i*phase)|11...1⟩)/√2
        
        Args:
            phase: The relative phase in radians
            measure: Whether to include measurement operations in the circuit
            
        Returns:
            A quantum circuit that prepares a GHZ state with the specified phase
        """
        # Create a new circuit with classical registers if needed
        qc = QuantumCircuit(self.num_qubits, len(self.qubit_indices) if measure else 0)
        
        # Apply Hadamard to the first qubit
        qc.h(self.qubit_indices[0])
        
        # Apply CNOT gates to entangle all qubits
        for i in range(len(self.qubit_indices) - 1):
            control = self.qubit_indices[i]
            target = self.qubit_indices[i + 1]
            qc.cx(control, target)
        
        # Apply phase to the last qubit
        qc.p(phase, self.qubit_indices[-1])
        
        # Add measurements if requested
        if measure:
            for i, qubit_idx in enumerate(self.qubit_indices):
                qc.measure(qubit_idx, i)
                
        self.circuit = qc
        return qc


if __name__ == "__main__":
    # Example usage
    ghz = GHZCircuit(num_qubits=5, qubit_indices=[0, 1, 2, 3, 4])
    circuit = ghz.generate_circuit()
    print(circuit.draw())
    
    # Example with phase
    ghz_phase = GHZCircuit(num_qubits=5)
    circuit_phase = ghz_phase.generate_circuit_with_phase(phase=3.14159/2)
    print(circuit_phase.draw())
